
from adt.tree import Tree
from compilers.basics import Token
from compilers.operator_based import FormalOperator, FormalOperand
from compilers.operator_based.calculator import Calculator

import re



class ExpressionLanguage:
    NUMBER = FormalOperand()
    PARENS = FormalOperand()
    LATEX = FormalOperand()
    ENCLOSED = FormalOperand()
    COMPOUND = FormalOperand()
    POW = FormalOperator(3, 1, 1)
    PLUS = FormalOperator(5, 1, 1, 0)
    MINUS = FormalOperator(5, 1, 1, 0)
    MUL = FormalOperator(4, 1, 1)
    

class SimplisticLexer(object):
    E = ExpressionLanguage
    
    RULES = {r"\^": E.POW,
             r"-": E.MINUS,
             r"\+": E.PLUS,
             r"\*": E.MUL,
             r"\d+": E.NUMBER,
             r"\s+": None,
             (r"\\([a-z])+({.*})*", "{}*"): E.LATEX,
             (r"\(.*\)", "()"): E.PARENS}
    
    def __init__(self):
        self.rules = [(self._matchexpr(x), y) for (x, y) in self.RULES.iteritems()] 
    
    def _matchexpr(self, expr_str):
        if isinstance(expr_str, str):
            return re.compile(expr_str)
        elif isinstance(expr_str, tuple):
            from compilers.operator_based import parens
            re_expr_str, bal_expr_str = expr_str
            re_expr = re.compile(re_expr_str)
            bal_expr = parens.BalancedExpression.compile(bal_expr_str)
            return parens.BothExpressions([re_expr, bal_expr])
        else:
            raise ValueError, expr_str
    
    def tokens(self, input):
        at = 0
        while input:
            tok = self.next_token(input)
            if tok is None:
                break
            elif tok.kind is not None:
                tok.at = at
                yield tok
            at += len(tok.text)
            input = input[len(tok.text):]
            
    
    def next_token(self, input):
        candidates = []
        
        # Try to match input according to the rules
        for regexp, kind in self.rules:
            mo = regexp.match(input)
            if mo is not None:
                if mo.groupdict():
                    tok = Token(kind, mo.groupdict())
                else:
                    tok = Token(kind, mo.group(0))
                candidates.append(tok)
            
        # Pick the longest sequence matched
        if not candidates:
            raise ValueError, "unrecognized token: " + input
        else:
            return max((len(tok.text), tok) for tok in candidates)[1]


class LaTeXElaboration(SimplisticLexer):
    RULES = {r"\\(?P<dir>[a-z]+){(?P<arg1>[^}]*)}{(?P<arg2>[^}]*)}": ExpressionLanguage.COMPOUND,
             r"\\(?P<dir>[a-z]+){(?P<arg>[^}]*)}": ExpressionLanguage.ENCLOSED }


    def __call__(self, expr_tree):
        if expr_tree.root.kind == ExpressionLanguage.LATEX:
            return self.make_tree(self.next_token(expr_tree.root.text))
        elif expr_tree.root.kind == ExpressionLanguage.PARENS:
            return self.sub_parser(expr_tree.root.text[1:-1])
        else:
            return Tree(expr_tree.root, [self(s) for s in expr_tree.subtrees])
        
    def make_tree(self, compound_token):
        tree = Tree(Token(compound_token.kind, compound_token.text['dir']))
        for sub in ['arg', 'arg1', 'arg2']:
            if sub in compound_token.text:
                tree.subtrees.append(self.sub_parser(compound_token.text[sub]))
        return tree


class LaTeXFormulaParser(object):
    
    def __init__(self):
        self.lex = SimplisticLexer()
        self.latex = LaTeXElaboration()
        self.latex.sub_parser = self
    
    def __call__(self, input):
        c = Calculator()
        c.process(self.lex.tokens(input))
        expr = c.result()
        expr = self.latex(expr)
        return expr



if __name__ == '__main__':
    inputs = [r"3^5",
              r"3^\textrm{5}",
              r"3^\frac{1}{2}",
              r"3^3 - 10",
              r"-3^3",
              r"-1+3^3",
              r"5-(3^4+1)"]
   
    fparser = LaTeXFormulaParser()
   
    for input in inputs:
        print input 
        print fparser(input)
        